package cn.itcast.star.graph.comfyui.client.config;

import cn.itcast.star.graph.comfyui.client.api.ComfyUI;
import cn.itcast.star.graph.comfyui.client.handler.ComfyUIMessageHandler;
import okhttp3.OkHttpClient;
import okhttp3.logging.HttpLoggingInterceptor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.web.socket.client.WebSocketClient;
import org.springframework.web.socket.client.WebSocketConnectionManager;
import org.springframework.web.socket.client.standard.StandardWebSocketClient;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

import java.util.concurrent.TimeUnit;

/**
 * ComfyUI配置类，用于创建ComfyUI相关的Bean实例
 */
@Configuration
@EnableWebSocket
public class ComfyUIConfig implements WebSocketConfigurer {

    /**
     * 创建ComfyUI Bean实例
     * 该方法配置了HTTP客户端、日志拦截器、连接超时等参数，并构建Retrofit实例
     * @return ComfyUI接口的实现实例
     */
    @Bean
    public ComfyUI comfyUI() {

        // 配置HTTP日志拦截器，用于记录请求和响应的详细信息
        HttpLoggingInterceptor loggingInterceptor = new HttpLoggingInterceptor();
        loggingInterceptor.setLevel(HttpLoggingInterceptor.Level.BODY);

        // 构建OkHttpClient，设置拦截器、重试机制和连接超时
        OkHttpClient okHttpClient = new OkHttpClient()
                .newBuilder()
                .addInterceptor(loggingInterceptor)
                .retryOnConnectionFailure(true)
                .connectTimeout(30, TimeUnit.SECONDS)
                .build();

        // 构建Retrofit实例，配置基础URL、HTTP客户端和JSON转换器
        Retrofit retrofit = new Retrofit.Builder()
                .baseUrl("http://127.0.0.1:8188")
                .client(okHttpClient)
                .addConverterFactory(JacksonConverterFactory.create())
                .build();
        return retrofit.create(ComfyUI.class);
    }

    /**
     * 创建ComfyUIMessageHandler Bean实例
     * 用于处理ComfyUI相关的消息
     * @return ComfyUIMessageHandler实例
     */
    @Bean
    public ComfyUIMessageHandler comfyUIMessageHandler() {
        return new ComfyUIMessageHandler();
    }

    /**
     * 创建WebSocket连接管理器Bean实例
     * 用于管理与ComfyUI的WebSocket连接
     * @param comfyUIMessageHandler 消息处理器实例
     * @return WebSocketConnectionManager实例
     */
    @Bean
    public WebSocketConnectionManager webSocketConnectionManager(ComfyUIMessageHandler comfyUIMessageHandler) {

        WebSocketClient webSocketClient = new StandardWebSocketClient();
        String url = "ws://127.0.0.1:8188/ws?clientId=star-graph";
        WebSocketConnectionManager connectionManager = new WebSocketConnectionManager(webSocketClient, comfyUIMessageHandler, url);
        // connectionManager.setAutoStartup(true); // 自动启动连接
        connectionManager.start();
        return connectionManager;
    }

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        // 注册WebSocket处理器，处理来自浏览器的WebSocket消息
        registry.addHandler(comfyUIMessageHandler(), "/ws")
                .setAllowedOrigins("*");// 允许跨域
    }
}

